import numpy as np
import matplotlib

matplotlib.use(backend="TkAgg")
import matplotlib.pyplot as plt
import pandas as pd
from math import isclose
import random
from collections import Counter
import numpy.random as npr
import numpy.linalg as linalg

np.set_printoptions(precision=4, suppress=True)


# compute stationary via power method
def stationary_power_method(P, tol=1e-10, max_iter=10000):
    n = P.shape[0] # 状态数量
    v = np.ones(n) / n  # 初始化均匀分布向量 [1/n, 1/n, ..., 1/n]
    for _ in range(max_iter):
        v_next = v.dot(P)
        # ||v_next - v||₁ = Σ|v_next[i] - v[i]|
        if linalg.norm(v_next - v, ord=1) < tol: # 检查收敛
            return v_next
        v = v_next
    return v

print("Section5: Stationary distribution computation & verification (irreducible finite chain)")
# create an irreducible apreriodic chain by mixing identity a bit
'''
这个转移矩阵描述了一个不可约且非周期的马尔可夫链：
所有状态互相可达（不可约）
没有周期性（对角元素都大于0）
这是一个"懒惰随机游走"的变体，每个状态都有较高的自环概率
'''
P5 = np.array([
    [0.5, 0.3, 0.2],
    [0.2, 0.5, 0.3],
    [0.3, 0.2, 0.5]
])
print(pd.DataFrame(P5))
s_power = stationary_power_method(P5)
'''
使用最小二乘法计算马尔可夫链平稳分布

平稳分布的基本方程
马尔可夫链的平稳分布 π 满足两个条件：
平稳性条件：π = πP
概率归一化条件：Σπ_i = 1
其中 π 是行向量，P 是转移矩阵。

步骤1：将平稳性条件改写为齐次方程
从 π = πP 可得：
π - πP = 0
π(I - P) = 0
其中 I 是单位矩阵。

步骤2：转置处理
为了便于求解，我们转置整个方程：
[π(I - P)]ᵀ = 0ᵀ
(I - P)ᵀπᵀ = 0
这里 πᵀ 是列向量。

步骤3：添加归一化条件
我们需要添加概率归一化条件 Σπ_i = 1，这可以写成：
[1, 1, ..., 1]πᵀ = 1

步骤4：组合成扩展系统
将两个条件组合成一个扩展的线性系统：

[A] [πᵀ] = [b]
[1ᵀ]      [1]
其中：
A = (I - P)ᵀ
1ᵀ = [1, 1, ..., 1]（全1行向量）
b = 0（零向量）
右侧扩展为 [0, 0, ..., 0, 1]ᵀ

步骤5：最小二乘公式
现在我们有系统：
Mx = c
其中：
M = [A; 1ᵀ]（A 下面加上全1行）
x = πᵀ
c = [0, 0, ..., 0, 1]ᵀ
这是一个超定系统（方程数多于未知数），我们使用最小二乘法求解：

minimize ‖Mx - c‖₂²
最小二乘解由正规方程给出：
MᵀMx = Mᵀc
解为：
x = (MᵀM)⁻¹Mᵀc



'''
A = P5.T - np.eye(3) # A = (p^T - I) = -(I-P)^T
A2 = np.vstack([A, np.ones(3)]) # M=[A; 1^T]
b2 = np.array([0,0,0,1.0]) # c=[0,0,0,1]^T
sol = linalg.lstsq(A2, b2, rcond=None)[0]
print("Stationary (power method):", s_power)
print("Stationary (lstsq):", sol)
print("Check sP - s (should be near 0):", sol.dot(P5) - sol)

